'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 5_with_langchain_summery_buffer.py
* @Time: 2025/10/24
* @All Rights Reserve By Brtc
'''
'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 4_with_conversation_tokens_buffer_memory.py
* @Time: 2025/10/24
* @All Rights Reserve By Brtc
'''
from operator import itemgetter

import dotenv
from langchain.memory import ConversationTokenBufferMemory, ConversationBufferWindowMemory, \
    ConversationSummaryBufferMemory
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from langchain_openai import ChatOpenAI

dotenv.load_dotenv()

#1、创建提示词
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是OpenAI 开发的机器人,请根据对应的上下文来回答用户的问题."),
    MessagesPlaceholder("history"),
    ("human", "{query}")
])
llm = ChatOpenAI(model="gpt-4o-mini")

moving_summary_buffer = ""
#记忆
memory = ConversationSummaryBufferMemory(max_token_limit=300,
                                         return_messages=True,
                                         input_key="query",
                                         llm = ChatOpenAI(model="gpt-4o-mini"))

#3、构建链应用
chain = (RunnablePassthrough.assign(history= RunnableLambda(memory.load_memory_variables)|itemgetter("history"))
         |prompt
         |llm
         |StrOutputParser())

#4、构建对话命令行
while True:
    query = input("Human:")
    if query == "exit":
        print("Goodbye")
        exit(0)

    chain_input = {"query": query, "history":[]}
    response = chain.stream(chain_input)
    print("AI:", flush=True, end = "")
    ai_content=""
    for chunk in response:
        print(chunk, flush=True, end="")
        ai_content += chunk
    memory.save_context(chain_input, {"output": ai_content})
    print("")
    print("history:", memory.load_memory_variables({}))















